// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::{
    bls12381::{BLS12381PrivateKey, BLS12381PublicKey, BLS12381Signature},
    ed25519::{Ed25519PrivateKey, Ed25519PublicKey, Ed25519Signature},
    traits::*,
    unit_tests::uniform_keypair_strategy,
};

use crypto::hash::HashValue;

use core::convert::TryFrom;
use crypto_derive::SilentDebug;
use failure::prelude::*;
use proptest::prelude::*;
use serde::{Deserialize, Serialize};

// Here we aim to make a point about how we can build an enum generically
// on top of a few choice signing scheme types. This enum implements the
// VerifyingKey, SigningKey for precisely the types selected for that enum
// (here, Ed25519(PublicKey|PrivateKey|Signature) and BLS(...)).
//
// Note that we do not break type safety (see towards the end), and that
// this enum can safely be put into the usual collection (Vec, HashMap).
//
// Note also that *nothing* in this definition requires knowing the details
// of the enum, so all of the below declarations can¹ be generated by a
// macro for any enum type making a coherent choice of concrete alternatives.
//
// ¹: and in the near future will << TODO(fga)

#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, Hash)]
enum PublicK {
    Ed(Ed25519PublicKey),
    BLS(BLS12381PublicKey),
}

#[derive(Serialize, Deserialize, SilentDebug)]
enum PrivateK {
    Ed(Ed25519PrivateKey),
    BLS(BLS12381PrivateKey),
}

#[allow(clippy::large_enum_variant)]
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum Sig {
    Ed(Ed25519Signature),
    BLS(BLS12381Signature),
}

impl From<&PrivateK> for PublicK {
    fn from(secret_key: &PrivateK) -> Self {
        match secret_key {
            PrivateK::Ed(pk) => PublicK::Ed(pk.into()),
            PrivateK::BLS(pk) => PublicK::BLS(pk.into()),
        }
    }
}

impl TryFrom<&[u8]> for PrivateK {
    type Error = CryptoMaterialError;
    fn try_from(bytes: &[u8]) -> std::result::Result<PrivateK, CryptoMaterialError> {
        Ed25519PrivateKey::try_from(bytes)
            .and_then(|ed_priv_key| Ok(PrivateK::Ed(ed_priv_key)))
            .or_else(|_err| {
                BLS12381PrivateKey::try_from(bytes)
                    .and_then(|bls_priv_key| Ok(PrivateK::BLS(bls_priv_key)))
            })
    }
}

impl ValidKey for PrivateK {
    fn to_bytes(&self) -> Vec<u8> {
        match self {
            PrivateK::BLS(privkey) => privkey.to_bytes().to_vec(),
            PrivateK::Ed(privkey) => privkey.to_bytes().to_vec(),
        }
    }
}

impl PublicKey for PublicK {
    type PrivateKeyMaterial = PrivateK;
    // TODO(fga): fix this!
    fn length() -> usize {
        std::cmp::max(BLS12381PublicKey::length(), Ed25519PublicKey::length())
    }
}

impl TryFrom<&[u8]> for PublicK {
    type Error = CryptoMaterialError;
    fn try_from(bytes: &[u8]) -> std::result::Result<PublicK, CryptoMaterialError> {
        Ed25519PublicKey::try_from(bytes)
            .and_then(|ed_priv_key| Ok(PublicK::Ed(ed_priv_key)))
            .or_else(|_err| {
                BLS12381PublicKey::try_from(bytes)
                    .and_then(|bls_priv_key| Ok(PublicK::BLS(bls_priv_key)))
            })
    }
}

impl ValidKey for PublicK {
    fn to_bytes(&self) -> Vec<u8> {
        match self {
            PublicK::BLS(pubkey) => pubkey.to_bytes().to_vec(),
            PublicK::Ed(pubkey) => pubkey.to_bytes().to_vec(),
        }
    }
}

impl PrivateKey for PrivateK {
    type PublicKeyMaterial = PublicK;
}

impl SigningKey for PrivateK {
    type VerifyingKeyMaterial = PublicK;
    type SignatureMaterial = Sig;

    fn sign_message(&self, message: &HashValue) -> Sig {
        match self {
            PrivateK::Ed(ed_priv) => Sig::Ed(ed_priv.sign_message(message)),
            PrivateK::BLS(bls_priv) => Sig::BLS(bls_priv.sign_message(message)),
        }
    }
}

impl Signature for Sig {
    type VerifyingKeyMaterial = PublicK;

    fn verify(&self, message: &HashValue, public_key: &PublicK) -> Result<()> {
        self.verify_arbitrary_msg(message.as_ref(), public_key)
    }

    fn verify_arbitrary_msg(&self, message: &[u8], public_key: &PublicK) -> Result<()> {
        match (self, public_key) {
            (Sig::Ed(ed_sig), PublicK::Ed(ed_pub)) => ed_sig.verify_arbitrary_msg(message, ed_pub),
            (Sig::BLS(bls_sig), PublicK::BLS(bls_pub)) => {
                bls_sig.verify_arbitrary_msg(message, bls_pub)
            }
            _ => bail!(
                "provided the wrong alternative in {:?}!",
                (self, public_key)
            ),
        }
    }

    fn to_bytes(&self) -> Vec<u8> {
        match self {
            Sig::Ed(sig) => sig.to_bytes().to_vec(),
            Sig::BLS(sig) => sig.to_bytes().to_vec(),
        }
    }
}

impl TryFrom<&[u8]> for Sig {
    type Error = CryptoMaterialError;
    fn try_from(bytes: &[u8]) -> std::result::Result<Sig, CryptoMaterialError> {
        Ed25519Signature::try_from(bytes)
            .and_then(|ed_sig| Ok(Sig::Ed(ed_sig)))
            .or_else(|_err| {
                BLS12381Signature::try_from(bytes).and_then(|bls_sig| Ok(Sig::BLS(bls_sig)))
            })
    }
}

impl VerifyingKey for PublicK {
    type SigningKeyMaterial = PrivateK;
    type SignatureMaterial = Sig;
}

///////////////////////////////////////////////////////
// End of declarations — let's now prove type safety //
///////////////////////////////////////////////////////
proptest! {
    #![proptest_config(ProptestConfig::with_cases(20))]

    #[test]
    fn test_keys_mix(
        hash in any::<HashValue>(),
        ed_keypair1 in uniform_keypair_strategy::<Ed25519PrivateKey, Ed25519PublicKey>(),
        ed_keypair2 in uniform_keypair_strategy::<Ed25519PrivateKey, Ed25519PublicKey>(),
        bls_keypair in uniform_keypair_strategy::<BLS12381PrivateKey, BLS12381PublicKey>()
    ) {
        // this is impossible to write statically, due to the trait not being
        // object-safe (voluntarily)
        // let mut l: Vec<Box<dyn PrivateKey>> = vec![];
        let mut l: Vec<Ed25519PrivateKey> = vec![];
        l.push(ed_keypair1.private_key);
        let ed_key = l.pop().unwrap();
        let signature = ed_key.sign_message(&hash);

        // This is business as usual
        prop_assert!(signature.verify(&hash, &ed_keypair1.public_key).is_ok());

        // This is impossible to write, and generates:
        // expected struct `ed25519::Ed25519PublicKey`, found struct `bls12381::BLS12381PublicKey`
        // prop_assert!(signature.verify(&hash, &bls_keypair.public_key).is_ok());

        let mut l2: Vec<PrivateK> = vec![];
        l2.push(PrivateK::BLS(bls_keypair.private_key));
        l2.push(PrivateK::Ed(ed_keypair2.private_key));

        let ed_key = l2.pop().unwrap();
        let ed_signature = ed_key.sign_message(&hash);

        // This is still business as usual
        let ed_pubkey2 = PublicK::Ed(ed_keypair2.public_key);
        let good_sigver = ed_signature.verify(&hash, &ed_pubkey2);
        prop_assert!(good_sigver.is_ok(), "{:?}", good_sigver);

        // but this still fails, as expected
        let bls_pubkey = PublicK::BLS(bls_keypair.public_key);
        let bad_sigver = ed_signature.verify(&hash, &bls_pubkey);
        prop_assert!(bad_sigver.is_err(), "{:?}", bad_sigver);

        // And now just in case we're confused again, we pop in the
        // reverse direction
        let bls_key = l2.pop().unwrap();
        let bls_signature = bls_key.sign_message(&hash);

        // This is still business as usual
        let good_sigver = bls_signature.verify(&hash, &bls_pubkey);
        prop_assert!(good_sigver.is_ok(), "{:?}", good_sigver);

        // but this still fails, as expected
        let bad_sigver = bls_signature.verify(&hash, &ed_pubkey2);
        prop_assert!(bad_sigver.is_err(), "{:?}", bad_sigver);
    }
}
